import pandas as pd
import numpy as np
# 读取CSV文件
import my_model
import torch
import argparse
from tqdm import tqdm, trange
import data
import torch.nn as nn
import matplotlib.pyplot as plt
# 获得训练集

criterion = nn.MSELoss()  # 使用均方误差损失函数计算MSE


def get_train_data(file_path, edge_pth):
    df = pd.read_csv(file_path, encoding='utf-8')
    edge_df = pd.read_csv(edge_pth, encoding='utf-8')
    df.head()

    # 地理位置哈希值
    geohasd_df_dict = {}
    # 日期的索引
    date_df_dict = {}
    number_hash = 0
    number_date = 0
    #这里是把数据中复杂的id（字符串）映射成简单的id（数值，递增）
    for i in df["geohash_id"]:
        if i not in geohasd_df_dict.keys():
            geohasd_df_dict[i] = number_hash
            number_hash += 1
    for i in df["date_id"]:
        if i not in date_df_dict.keys():
            date_df_dict[i] = number_date
            number_date += 1

    # 构建一个三维的数组
    new_data = np.zeros(
        (len(date_df_dict), len(geohasd_df_dict), 38))  # [len(geohasd_df_dict) * [[0]]] * len(date_df_dict)
    print('ganfeini ')
    # 遍历原始的数据，
    for index, row in df.iterrows():
        # print(index)
        #这里是将geohash_id和date_id的索引值存下来，index是行索引，row表示当前行的数据
        hash_index, date_index = geohasd_df_dict[row["geohash_id"]], date_df_dict[row["date_id"]]
        # print(hash_index)
        # print(date_index)
        # exit()
        # 将时间index加到里面
        # print(len(np.array([date_index]+list(row.iloc[2:]))))
        # 用两个索引构成的位置信息指示当前数据，当前节点所在位置，这里对两个列表合在一起
        new_data[date_index][hash_index] = np.array([date_index] + list(row.iloc[2:]))
    # new_data = np.array(new_data)
    # x_train,y_train = new_data[:, :-2], new_data[:, -2:]
    # print(len(geohasd_df_dict))
    # exit()
    # print(x_train.shape)
    # print(y_train.shape)
    # 这里构建邻接矩阵其中mask表示1为有边，0无边， value_mask表示有值
    # 并且这里我考虑mask是一个无向图，如果有向删除x_mask[date_index][point2_index][point1_index],value_mask同理
    # x_mask表示一个无向图，x_edge_df存储边的特征
    x_mask = np.zeros((len(date_df_dict), len(geohasd_df_dict), len(geohasd_df_dict), 1), dtype=float)
    x_edge_df = np.zeros((len(date_df_dict), len(geohasd_df_dict), len(geohasd_df_dict), 2), dtype=float)

    for index, row in edge_df.iterrows():
        # print(index)
        # 如果存在边但不存在节点就舍弃
        if row["geohash6_point1"] not in geohasd_df_dict.keys() or row["geohash6_point2"] not in geohasd_df_dict.keys():
            continue
        point1_index, point2_index, F_1, F_2, date_index = geohasd_df_dict[row["geohash6_point1"]], geohasd_df_dict[
            row["geohash6_point2"]] \
            , row["F_1"], row["F_2"], date_df_dict[row["date_id"]]
        x_mask[date_index][point1_index][point2_index] = 1
        x_mask[date_index][point2_index][point1_index] = 1
        x_edge_df[date_index][point1_index][point2_index] = [F_1, F_2]
        x_edge_df[date_index][point2_index][point1_index] = [F_1, F_2]
    # print(data)

    return geohasd_df_dict, date_df_dict, new_data, x_mask, x_edge_df


def eval(model, dataset, args):
    model.eval()
    with torch.no_grad():
        dev_loss = 0.0
        for j in trange(dataset.batch_count):
            x_date, x_feature, x_mask_data, x_edge_data, x_tags = dataset.get_batch(j)
            act_pre, con_pre = model(x_date, x_feature, x_mask_data)
            predict = torch.cat((act_pre, con_pre), dim=-1)
            loss = criterion(predict, x_tags)
            dev_loss += loss
        print("this epoch dev loss is {}".format(dev_loss))
        model.train()


def train(args):
    geohasd_df_dict, date_df_dict, x_train, x_mask, x_edge_df = get_train_data('./train_90.csv',
                                                                               "./edge_90.csv")
    # 分割各种训练集测试集
    x_train, x_dev = torch.tensor(x_train[:int(len(x_train) * args.rat)]), torch.tensor(
        x_train[int(len(x_train) * args.rat):])
    x_mask_train, x_mask_dev = torch.tensor(x_mask[:int(len(x_mask) * args.rat)]), torch.tensor(
        x_mask[int(len(x_mask) * args.rat):])
    x_edge_train, x_edge_dev = torch.tensor(x_edge_df[:int(len(x_edge_df) * args.rat)]), torch.tensor(
        x_edge_df[int(len(x_edge_df) * args.rat):])

    date_emb = 5
    # 这里的x包含了date_id+F35个特征+2个y值的
    # train_activate = torch.tensor(y_train[:, -2])
    # train_consume = torch.tensor(y_train[:, -1])

    # rmse_loss = torch.sqrt(mse_loss)
    model = my_model.GAT(date_emb=[len(date_df_dict), date_emb], nfeat=35, nhid=64, dropout=0.3, alpha=0.3,nheads=8).to(args.device)
    # model = my_model.BILSTM(date_emb =[len(date_df_dict),date_emb], nfeat=35, nhid=64, dropout=0.3, alpha=0.3, nheads=8).to(args.device)
    optimizer = torch.optim.Adam(params=model.parameters(), lr=args.lr,weight_decay=0.01)
    # optimizer = torch.optim.AdamW(params=model.parameters(), lr=args.lr,weight_decay=0.01)
    # 学习率缩放
    # scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.decline, gamma=0.5, last_epoch=-1)
    model.train()
    trainset = data.DataIterator(x_train, x_mask_train, x_edge_train, args)
    valset = data.DataIterator(x_dev, x_mask_dev, x_edge_dev, args)

    # 记录损失函数
    trainloss=[]
    for indx in range(args.epochs):
        train_all_loss = 0.0
        for j in trange(trainset.batch_count):
            x_date, x_feature, x_mask_data, x_edge_data, x_tags = trainset.get_batch(j)
            act_pre, con_pre = model(x_date, x_feature, x_mask_data)
            predict = torch.cat((act_pre, con_pre), dim=-1)

            loss = criterion(predict, x_tags)
            train_all_loss += loss
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        print(f'epoch is {indx}')
        print('this epoch train loss :{0}'.format(train_all_loss))
        trainloss.append(train_all_loss)
        # scheduler.step()
        eval(model, valset, args)

    torch.save(model, 'model.pth')






def add_columns_with_zeros(input_file, output_file):
    # 读取 CSV 文件
    df = pd.read_csv(input_file)

    # 添加两列，并将数据填充为 0
    df['active_index'] = 0
    df['consume_index'] = 0

    # 保存到新的 CSV 文件
    df.to_csv(output_file, index=False)


# 在A榜测试集测试
def evaluate_on_test(model, dataset, args):
    model.eval()
    predictions = []

    with torch.no_grad():
        for j in trange(dataset.batch_count):
            x_date, x_feature, x_mask_data, x_edge_data, _ = dataset.get_batch(j)
            act_pre, con_pre = model(x_date, x_feature, x_mask_data)
            predict = torch.cat((act_pre, con_pre), dim=-1)

            # 使用 torch.meshgrid 生成类似 ndindex 的索引
            date_idx, geohash_idx = torch.meshgrid(torch.arange(predict.shape[0]), torch.arange(predict.shape[1]))

            # 这里可以根据你的需要对预测结果进行后处理
            # 例如反归一化等
            # Extract geohash_id, consumption_level, activity_level, date_id
            for date_idx_val, geohash_idx_val in zip(date_idx.flatten(), geohash_idx.flatten()):
                geohash_id = list(geohasd_df_dict.keys())[list(geohasd_df_dict.values()).index(geohash_idx_val)]
                date_id = list(date_df_dict.keys())[list(date_df_dict.values()).index(date_idx_val)]

                # 合并为一个字符串
                prediction_str = f'{geohash_id}\t{predict[date_idx_val, geohash_idx_val, -2].item()}\t' \
                                 f'{predict[date_idx_val, geohash_idx_val, -1].item()}\t{date_id}'

                predictions.append(prediction_str)

    # 创建一个 DataFrame 保存预测结果
    predictions_df = pd.DataFrame(predictions, columns=['geohash_id	consumption_level	activity_level	date_id'])

    # 保存为 CSV 文件
    predictions_df.to_csv('model_predictions_test.csv', index=False)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', type=int, default=333,
                        help='training epoch number')
    parser.add_argument('--batch_size', type=int, default=4,
                        help='batch_size')
    parser.add_argument('--device', type=str, default="cuda",
                        help='gpu or cpu')
    parser.add_argument('--lr', type=float, default=1e-2,
                        )
    parser.add_argument('--rat', type=float, default=0.9, )

    parser.add_argument('--decline', type=int, default=30, help="number of epochs to decline")

    #在测试集添加两列
    #add_columns_with_zeros('eval/node_test_4_A.csv', 'eval/node.csv')

    train(parser.parse_args())
    model=torch.load('model.pth')
    geohasd_df_dict, date_df_dict, x_test, x_mask_test, x_edge_test = get_train_data('eval/node.csv','eval'
                                                                                         '/edge_test_4_A.csv')
    x_test = torch.tensor(x_test)
    x_mask_test = torch.tensor(x_mask_test)
    x_edge_test = torch.tensor(x_edge_test)
    testset = data.DataIterator(x_test, x_mask_test, x_edge_test, parser.parse_args())
    evaluate_on_test(model,testset,parser.parse_args())

    # %%

# gcn = GCN(n_features=51, hidden_dim=64, dropout=0.3, n_classes=2)

# %%

# date_node = np.concatenate([sample.iloc[:, 2:37].values, date_embed], axis=1)
# date_node.shape
#
# # %%
#
# adj = torch.randn(90, 90)
# x_node = torch.from_numpy(date_node).type_as(adj)
# predict = gcn(x_node, adj)
# predict.shape, predict
#
# # %%
#
# gat = GAT(nfeat=51, nhid=64, nclass=2, dropout=0.3, alpha=0.3, nheads=8)
#
# # %%
#
# adj = torch.randn(90, 90)
# x_node = torch.from_numpy(date_node).type_as(adj)
# predict = gat(x_node, adj)
# predict.shape, predict
#
# # %%
#
# actual_values = sample.iloc[:, 37:]  # active_index，consume_index
#
# # %%
#
# import torch
# import torch.nn as nn
#
# # 定义预测值和实际观测值
# actual_values = torch.from_numpy(sample.iloc[:, 37:].values)  # 实际观测值
#
# # 计算RMSE损失
# criterion = nn.MSELoss()  # 使用均方误差损失函数计算MSE
# mse_loss = criterion(predict, actual_values)
# rmse_loss = torch.sqrt(mse_loss)
#
# # 打印RMSE损失
# print("RMSE Loss:", rmse_loss.item())
#
# # %%
#
# x_node, adj
#
# # %%
#
# import pandas as pd
#
# # 读取CSV文件
# edge = pd.read_csv('/home/cike/workspace/data/datamining/edge_90.csv')
#
# # %%
#
# edge.head()
#
# # %%
#
# geohash_id = df.geohash_id
#
# # %%
#
# geohash_id.head()
#
# # %%
#
# align = edge[(edge['geohash6_point1'] == df.geohash_id[0]) & (edge['date_id'] == df.date_id[0])]
#
# # %%
#
# align
#
# # %%
#
# df.geohash_id.drop_duplicates()
